/******************************************************************************
 * Copyright 2022 The Untouch Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/
#include "modules/calibration/lidar_imu_calibrator/include/self_lidar_imu_calibrator.h"

#include <cmath>
#include <future>
#include <iomanip>
#include <unistd.h>

#include "nlopt.hpp"
#include "pcl/io/pcd_io.h"
#include "pcl/common/transforms.h"

#include "opencv2/opencv.hpp"

namespace lidar_imu_calibrator {

SelfLidarImuCalibrator::SelfLidarImuCalibrator() {
  select_frames_.clear();
  select_pairs_positive_.clear();
  select_paris_reverse_.clear();
};

double SelfLidarImuCalibrator::Calibrate(
    std::vector<OdometryFrame::Ptr>& frames,
    Eigen::Isometry3d& extrinsics) {  

  if (frames.size() < 10) {
    AERROR << "[SelfLidarImuCalibrator::SelectFrames]: too less lidar frames! ";
    return false;
  }
  if (int(select_paris_reverse_.size()) < max_bin_size_){
    if (!SelectFrames(frames)) 
      return false;
  }
  
  if (!IsCalibrate()) return false;
  return StarCalibration(select_frames_, extrinsics);
}

bool SelfLidarImuCalibrator::SaveCloud(const std::string &path, const Eigen::Isometry3d& extrinsics){
  PointCloud merged_cloud_1, merged_cloud_2;
  GetTwoCombinedPointcloud(extrinsics, &merged_cloud_1, &merged_cloud_2);

  pcl::PCDWriter pcd_writer;
  if (!merged_cloud_1.empty())
    pcd_writer.writeBinaryCompressed(path, merged_cloud_1);
  return true;
}

bool SelfLidarImuCalibrator::SelectFrames(
    const std::vector<OdometryFrame::Ptr>& frames) {
  draw_board_ = cv::Mat(1080, 1080, CV_8UC3, cv::Scalar(0, 0, 0));
  double cloud_time_, cumulative = 0, pre_select_time_ = 0;
  std::vector<std::pair<double, double>> positions;
  std::vector<double> euler_yaw_pre_imu_;
  for (size_t i = 0; i < frames.size(); i++) {
    // std::cout << "[SelectFrames]: loop " << i << std::endl;
    cloud_time_ = frames[i]->GetFrameTimestamp();
    if (pre_select_time_ == 0) pre_select_time_ = cloud_time_;

    EulerAnglesZXYd euler(
        Eigen::Quaterniond(frames[i]->GetFramePose().linear()));
    double yaw = euler.yaw() * 180.0 / M_PI;
    if (yaw < 0) yaw += 360;

    euler_yaw_pre_imu_.push_back(yaw);


    if (!star_select_frames_){
      if (IsStarSelect(euler_yaw_pre_imu_)){
        euler_yaw_pre_imu_.clear();
        positions.clear();
        base_pose_ = frames[i + 1]->GetImuFramePose();
      }
      continue;
    }
    positions.emplace_back(frames[i]->GetFramePose().translation().x() - base_pose_.translation().x(),
                           frames[i]->GetFramePose().translation().y() - base_pose_.translation().y());
    select_frames_.push_back(frames[i]);

    DrawPose(base_pose_.inverse() * frames[i]->GetFramePose(), 2, cv::Scalar(255,255,255));
    if (cumulative < max_angle_cumulative_) {
      static double yaw_star = yaw;
      float bias = fabs(yaw - yaw_star);
      cumulative += bias > 60 ? 0 : bias;
      yaw_star = yaw;
      first_circle_idex_ = int(positions.size());
      continue;
    }

    if (fabs(cloud_time_ - pre_select_time_) > time_limit_) {
      int match = PoseSelection(positions, euler_yaw_pre_imu_);
      if (match > 0){
        DrawPose(base_pose_.inverse() * (*select_frames_.rbegin())->GetFramePose(), 4, cv::Scalar(0,0,255));
        DrawPose(base_pose_.inverse() * select_frames_[match]->GetFramePose(), 4, cv::Scalar(0,255,0));
      }
      pre_select_time_ = cloud_time_;
    }
    if (int(select_pairs_positive_.size()) >= max_bin_size_) break;
  }
  // cv::destroyAllWindows();
  return true;
}

void SelfLidarImuCalibrator::DrawPose(const Eigen::Isometry3d& pose, int mark_radius, const cv::Scalar& color){
  static const double kDrawBoardRange = 50;
  int width = draw_board_.cols;
  int height = draw_board_.rows;

  int x =
      static_cast<int>(pose.translation().x() / kDrawBoardRange * width) +
      width / 2;
  int y =
      static_cast<int>(pose.translation().y() / kDrawBoardRange * height) +
      height / 2;

  if (x >= 0 && x < width && y >= 0 && y < height) {
    cv::circle(draw_board_, cv::Point(x, y), mark_radius, color);
  }

  cv::imshow("Draw Board", draw_board_);
  cv::waitKey(100);
}

int SelfLidarImuCalibrator::PoseSelection(
    const std::vector<std::pair<double, double>>& positions,
    const std::vector<double>& euler_yaw_pre_imu) {
  if (positions.size() < 2 || positions.size() != euler_yaw_pre_imu.size())
    return -1;

  double current_x = positions.rbegin()->first,
         current_y = positions.rbegin()->second;
  double current_yaw = *euler_yaw_pre_imu.rbegin();
  double min_dist = min_dist_;
  int bin_id = -1, matched_index = -1;

  for (int i = int(positions.size()) - first_circle_idex_; i > 0; i--) {
    double dist = sqrt(pow((positions[i].first - current_x), 2) +
                       pow((positions[i].second - current_y), 2));
    if (dist > min_dist) continue;

    double yaw = current_yaw - euler_yaw_pre_imu[i];
    if (yaw < 0) yaw += 360.0;

    if (yaw < bound_forward_ || yaw > (360.0 - bound_forward_)) {
      bin_id = 0;
    }
    if (yaw < (180.0 + bound_forward_) && yaw > (180.0 - bound_forward_)) {
      bin_id = 1;
    }
    if ((1 == bin_id && int(select_pairs_positive_.size()) < max_bin_size_) ||
        (0 == bin_id && int(select_paris_reverse_.size()) < max_bin_size_)) {
      min_dist = dist;
      matched_index = i;
    }
  }
  if (bin_id == 0 && int(select_pairs_positive_.size()) < max_bin_size_) {
    select_pairs_positive_.emplace_back(matched_index, positions.size() - 1);
  } else if (bin_id == 1 && int(select_paris_reverse_.size()) < max_bin_size_) {
    select_paris_reverse_.emplace_back(matched_index, positions.size() - 1);
  }
  
  return matched_index;
}

bool SelfLidarImuCalibrator::IsStarSelect(const std::vector<double> &euler_yaw_pre_imu){
  if (euler_yaw_pre_imu.size() < 30) return false;
  double max_bias = 60;
  for (size_t i = euler_yaw_pre_imu.size() - 1; i > euler_yaw_pre_imu.size() - 11; i--)
  {
    if (fabs(euler_yaw_pre_imu[i] - euler_yaw_pre_imu[i-1]) > max_bias)
      return false;
    /* code */
  }
  star_select_frames_ = true;
  return true;
}

bool SelfLidarImuCalibrator::IsCalibrate(){
  if (int(select_pairs_positive_.size()) >= max_bin_size_) return true;
  if (int(select_pairs_positive_.size()) >= max_bin_size_ && 
      int(select_paris_reverse_.size()) >= max_bin_size_)
      return true;
  return false;
}

double SelfLidarImuCalibrator::StarCalibration(const std::vector<OdometryFrame::Ptr>& frames, 
                        Eigen::Isometry3d& extrinsics){
  EulerAnglesZXYd extrin_euler( Eigen::Quaterniond(extrinsics.linear()));
  std::vector<double> extrin_euler_lidar_2_imu_(6, 0.0);
  extrin_euler_lidar_2_imu_ = {
      extrinsics.translation().x(), extrinsics.translation().y(), extrinsics.translation().z(),
      // 0, 90, 0 };
            extrin_euler.roll() * 180 / M_PI, extrin_euler.pitch() * 180 / M_PI, extrin_euler.yaw() * 180 / M_PI };
  
  std::vector<double> lb = {-translation_range_, -translation_range_,
                            -translation_range_, -angular_range_,
                            -angular_range_,     -angular_range_};
  std::vector<double> ub = {translation_range_, translation_range_,
                            translation_range_, angular_range_,
                            angular_range_,     angular_range_};
  
  double roll = extrin_euler_lidar_2_imu_[3], pitch = extrin_euler_lidar_2_imu_[4], yaw = extrin_euler_lidar_2_imu_[5];
  for (size_t i = 0; i < 6; ++i) {
    lb[i] += extrin_euler_lidar_2_imu_[i];
    ub[i] += extrin_euler_lidar_2_imu_[i];
  }

  if (time_cal_) {
    extrin_euler_lidar_2_imu_.push_back(0.0);
    ub.push_back(max_time_offset_);
    lb.push_back(-max_time_offset_);
  }

  double error = Optimize(lb, ub, this, &extrin_euler_lidar_2_imu_);

  Eigen::Matrix<double, 6, 1> vec;
  for (size_t i = 0; i < 6; ++i) {
    vec[i] = extrin_euler_lidar_2_imu_[i];
  }
  std::cout << "cali first : " << roll << " " << pitch << " " << yaw << std::endl;
  // std::cout << "tran first : " << extrinsics.translation().x() << " " << extrinsics.translation().y() << " " << extrinsics.translation().z() << std::endl;
  std::cout << "cali result: " << extrin_euler_lidar_2_imu_[3] << " " << extrin_euler_lidar_2_imu_[4] << " " << extrin_euler_lidar_2_imu_[5] << std::endl;
  // std::cout << "tran result: " << extrin_euler_lidar_2_imu_[0] << " " << extrin_euler_lidar_2_imu_[1] << " " << extrin_euler_lidar_2_imu_[2] << std::endl;
  Euler2Isometry3d(vec, extrinsics);
  // std::cout << Eigen::Quaterniond(extrinsics.linear()).coeffs() << std::endl;
  return error;
}

void SelfLidarImuCalibrator::SetTransformLidar2Imu(const Eigen::Matrix<double, 6, 1> &vec){
  Euler2Isometry3d(vec, T_l_o_);
}

void SelfLidarImuCalibrator::Euler2Isometry3d(const Eigen::Matrix<double, 6, 1> &vec, Eigen::Isometry3d& extrinsics){
  Eigen::Translation3d translation(vec[0], vec[1], vec[2]);
  EulerAnglesZXYd euler_angle(vec[3] * M_PI / 180, vec[4] * M_PI / 180,
                              vec[5] * M_PI / 180);
  Eigen::Quaterniond rotation = euler_angle.ToQuaternion();

  extrinsics = translation * rotation;
}

void SelfLidarImuCalibrator::GetTwoCombinedPointcloud(
    const Eigen::Isometry3d& extrinsics, PointCloud* merged_cloud_1,
    PointCloud* merged_cloud_2) {
  merged_cloud_1->clear();
  merged_cloud_2->clear();
  pcl::PCDWriter pcd_writer;
  for (size_t i = 0; i < select_pairs_positive_.size(); i++)
  {
    auto p = select_pairs_positive_[i];
    Eigen::Isometry3d lidar_pose = base_pose_.inverse() * select_frames_[p.second]->GetImuFramePose();

    lidar_pose = lidar_pose * extrinsics;
    // Eigen::Isometry3d lidar_pose2 = select_frames_[p.second]->GetFramePose();

    PointCloudPtr cloud(new PointCloud);
    select_frames_[p.second]->GetPointCloud(cloud);

    if (!cloud->empty()){
      pcl::transformPointCloud(*cloud, *cloud, lidar_pose.cast<float>().matrix());
      *merged_cloud_1 += *cloud;
    }

    lidar_pose = base_pose_.inverse() * select_frames_[p.first]->GetImuFramePose();
    lidar_pose = lidar_pose * extrinsics;
    select_frames_[p.first]->GetPointCloud(cloud);
    if (!cloud->empty()){
      pcl::transformPointCloud(*cloud, *cloud, lidar_pose.cast<float>().matrix());
      *merged_cloud_1 += *cloud;
    }
  }

  std::string savePath1 = "/apollo/data/lidar_imu/positive.pcd";
  if (!merged_cloud_1->empty() && access((char *)savePath1.c_str(), F_OK) == -1){
    pcd_writer.writeBinaryCompressed(savePath1, *merged_cloud_1);
  }
  // for (size_t i = 0; i < select_paris_reverse_.size(); i++)
  // {
  //   auto p = select_paris_reverse_[i];
  //   Eigen::Isometry3d lidar_pose = base_pose_.inverse() * select_frames_[p.second]->GetFramePose() * extrinsics;
  //   PointCloudPtr cloud(new PointCloud);
  //   select_frames_[p.second]->GetPointCloud(cloud);
  //   pcl::transformPointCloud(*cloud, *cloud, lidar_pose.matrix());
  //   (*merged_cloud_2) += (*cloud);

  //   lidar_pose = base_pose_.inverse() * select_frames_[p.first]->GetFramePose() * extrinsics;
  //   select_frames_[p.first]->GetPointCloud(cloud);
  //   pcl::transformPointCloud(*cloud, *cloud, lidar_pose.matrix());
  //   (*merged_cloud_2) += (*cloud);
  // }
  // if (!merged_cloud_2->empty()){
  //   std::string savePath2 = "/apollo/data/lidar_imu/reverse.pcd";
  //   pcd_writer.writeBinaryCompressed(savePath2, *merged_cloud_2);
  // }
}

float SelfLidarImuCalibrator::lidarOdomKNNError(){
  PointCloud merged_cloud_1, merged_cloud_2;
  GetTwoCombinedPointcloud(T_l_o_, &merged_cloud_1, &merged_cloud_2);

  // 也可以尝试整个点云对整个点云的误差，就跟lidar_align一样？
  return lidarOdomKNNError(merged_cloud_1, merged_cloud_1);
}

float SelfLidarImuCalibrator::lidarOdomKNNError(
    const PointCloud& base_pointcloud, const PointCloud& combined_pointcloud) {
  // shared_pointer needed by kdtree, no-op destructor to prevent it trying to
  // clean it up after use
  PointCloud::ConstPtr combined_pointcloud_ptr(&combined_pointcloud,
                                               [](const PointCloud*) {});

  pcl::KdTreeFLANN<Point> kdtree;
  kdtree.setInputCloud(combined_pointcloud_ptr);

  float max_dist = local_knn_max_dist_;

  size_t k = knn_k_;
  // if searching own cloud add one to k as a point will always match to itself
  if (&base_pointcloud == &combined_pointcloud) {
    ++k;
  }

  // small amount of threading here to take edge off of bottleneck break knn
  // lookup up into several smaller problems each running in their own thread
  std::vector<std::future<float>> errors;
  for (size_t start_idx = 0; start_idx < base_pointcloud.size();
       start_idx += knn_batch_size_) {
    size_t end_idx = start_idx + std::min(base_pointcloud.size() - start_idx,
                                          static_cast<size_t>(knn_batch_size_));
    errors.emplace_back(
        std::async(std::launch::async, SelfLidarImuCalibrator::kNNError,
                   kdtree, base_pointcloud, k, max_dist, start_idx, end_idx));
  }

  // wait for threads to finish and grab results
  float total_error = 0.0f;
  for (std::future<float>& error : errors) {
    total_error += error.get();
  }

  return total_error;
}

float SelfLidarImuCalibrator::kNNError(const pcl::KdTreeFLANN<Point>& kdtree,
                                       const PointCloud& pointcloud,
                                       const size_t k, const float max_dist,
                                       const size_t start_idx,
                                       const size_t end_idx) {
  std::vector<int> kdtree_idx(k);
  std::vector<float> kdtree_dist(k);

  float error = 0;

  for (size_t idx = start_idx; idx < std::min(pointcloud.size(), end_idx);
       ++idx) {
    kdtree.nearestKSearch(pointcloud[idx], k, kdtree_idx, kdtree_dist);
    for (const float& x : kdtree_dist) {
      error += std::min(x, max_dist);
    }
  }

  return error;
}

double SelfLidarImuCalibrator::LidarOdomMinimizer(const std::vector<double>& x,
                                                  std::vector<double>& grad,
                                                  void* f_data) {
  SelfLidarImuCalibrator* od = static_cast<SelfLidarImuCalibrator*>(f_data);
  Eigen::Matrix<double, 6, 1> vec;
  vec.setZero();
  for (size_t i = 0; i < 6; ++i) {
    vec[i] = x[i];
  }

  if (x.size() > 6) {
    od->SetTransformLidar2Imu(vec);
  }
  od->SetTransformLidar2Imu(vec);

  double error = od->lidarOdomKNNError();

  static int i = 0;

  std::cout << std::fixed << std::setprecision(4);
  std::cout << " \e[1mx:\e[0m " << std::setw(6) << vec[0];
  std::cout << " \e[1my:\e[0m " << std::setw(6) << vec[1];
  std::cout << " \e[1mz:\e[0m " << std::setw(6) << vec[2];
  std::cout << " \e[1mrx:\e[0m " << std::setw(6) << vec[3];
  std::cout << " \e[1mry:\e[0m " << std::setw(6) << vec[4];
  std::cout << " \e[1mrz:\e[0m " << std::setw(6) << vec[5];
  if (x.size() > 6) {
    std::cout << " \e[1mtime:\e[0m " << std::setw(6) << x[6];
  }
  std::cout << " \e[1mError:\e[0m " << std::setw(10) << error;
  std::cout << " \e[1mIteration:\e[0m " << i++ << '\r' << std::flush;

  return error;

}

double SelfLidarImuCalibrator::Optimize(const std::vector<double>& lb,
                                    const std::vector<double>& ub,
                                    SelfLidarImuCalibrator* data,
                                    std::vector<double>* x) {
  nlopt::opt opt;
  opt = nlopt::opt(nlopt::LN_BOBYQA, x->size());
  opt.set_lower_bounds(lb);
  opt.set_upper_bounds(ub);
  opt.set_maxeval(max_evals_);
  opt.set_xtol_abs(xtol_);

  opt.set_min_objective(LidarOdomMinimizer, data);

  double minf;
  nlopt::result result = opt.optimize(*x, minf);

  if (result == 1) {
    AINFO << "Optmize success!";
    return minf;
  }
  return -1;
}
}  // namespace lidar_imu_calibrator
